#pragma once
#include "wavAudio.h"

void wavAudio::readInfo(const char* filepath)
{
	LOG("read wav file info\n");

	File audio = SD.open(filepath);
	if (!audio)
	{
		LOG("Failed to open WAV file", filepath);
		while (1)
			;
	}
	LOG("WAV file opened, begin to read header\n");

	wav_header_t header;
	uint8_t buffer[sizeof(header)];
	audio.read(buffer, sizeof(buffer));
	memcpy(&header, buffer, sizeof(header));
	audio.close();

	this->format = header.audioFormat;
	this->channels = header.numChannels;
	this->sample_rate = header.sampleRate;
	this->bits_per_sample = header.bitsPerSample;
	this->data_length = header.byteRate * this->bits_per_sample / 16;
	this->samples_num = this->data_length * 8 / this->bits_per_sample;
	LOG("format: ", this->format, "\n");
	LOG("channels: ", this->channels, "\n");
	LOG("sample_rate: ", this->sample_rate, "\n");
	LOG("bits_per_sample: ", this->bits_per_sample, "\n");
	LOG("data_length: ", this->data_length, "\n");
	LOG("samples_num: ", this->samples_num, "\n");
	printMemoryInfo("read wav file header after");


	this->data = Malloc(int16_t, this->samples_num);
	printMemoryInfo("Allocate memory for data");

	LOG("begin to read data\n");
	audio = audio = SD.open(filepath); // 重新打开文件
	audio.seek(46); // 跳过头部信息

	int k = 0;
	while (audio.available() && k < this->samples_num) {
		// int16_t sample = audio.read() | (audio.read() << 8);
		// this->data[k++] = sample;
        uint8_t lowByte = audio.read();
        uint8_t highByte = audio.read();
        int16_t sample = (highByte << 8) | lowByte;  // 按小端序读取数据
        this->data[k++] = sample;
	}
	audio.close();
	printMemoryInfo("read wav file data after");
	LOG("read wav file info done\n");
}

bool wavAudio::processdata()
{
	LOG("\n\nbegin to process data\n");

	this->process_data = Malloc(int16_t, this->samples_num);
	printMemoryInfo("Allocate memory for process_data");

	this->max_data = std::numeric_limits<int16_t>::min();

	for (int i = 0; i < this->samples_num; i++)
	{
		int abs_data = this->data[i] > 0 ? this->data[i] : -this->data[i]; // 使用绝对值
		this->max_data = this->max_data > abs_data ? this->max_data : abs_data; // 更新最大值
	}
	LOG("max_data: ", this->max_data, "\n");

	// 确保max_data不为0，以避免后续计算时的除零错误
	if (this->max_data == 0)
	{
		LOG("All data points are zero, cannot process data.\n");
		for (int i = 0; i < this->samples_num; i++)
		{
			this->process_data[i] = 0;
		}
		return false;
	}

	this->_per = int(this->max_data * 0.1);
	LOG("_per: ", this->_per, "\n");

	int index = 0; // record the third num(not 0) position
	bool begin_clock_set = false;

	for (int i = 0; i < this->samples_num; i++)
	{
		// 使用绝对值比较
		int abs_data = this->data[i] > 0 ? this->data[i] : -this->data[i];
		if (abs_data <= this->_per)
		{
			this->process_data[i] = 0;
		}
		else
		{
			this->process_data[i] = int(this->data[i] * PEAK / this->max_data);
			index++;
			if (index == 4 && !begin_clock_set)
			{
				this->begin_clock = i;
				begin_clock_set = true;
			}
		}
	}

	index = 0; // record the end position
	bool end_clock_set = false;

	for (int i = this->samples_num - 1; i >= 0; i--)
	{
		if (this->process_data[i] != 0)
		{
			index++;
			if (index == 4 && !end_clock_set)
			{
				this->end_clock = i;
				end_clock_set = true;
				break;
			}
		}
	}

	LOG("begin_clock: ", this->begin_clock, "\n");
	LOG("end_clock: ", this->end_clock, "\n");
	LOG("process data done\n");
	printMemoryInfo("process data after");
	return true;
}

void wavAudio::calculatePitch()
{
    LOG("\n\nbegin to calculate pitch\n");
    
    int index[100] = {0};
    for (int i = this->begin_clock; i < this->end_clock; i++)
    {
        if (this->process_data[i] > 0 && this->process_data[i] < PEAK)
        {
            index[this->process_data[i] / 100 - 1]++;
        }
    }

    int n = 0;
    int m = INT_MIN;
    for (int i = 0; i < 100; i++)
    {
        if (index[i] > m)
        {
            n = i;
            m = index[i];
        }
    }

    if (n != 0)
    {
        this->pitch = n * 100 + 50;
    }
    else
    {
        this->pitch = 0;
    }

    LOG("pitch: ", this->pitch, "\n");
    printMemoryInfo("calculate pitch after");
}

void wavAudio::calculateAverageLoudness()
{
	LOG("\n\nbegin to calculate average loudness\n");
	int num = 0;
	long long sum = 0;
	for (int i = this->begin_clock; i < this->end_clock; i++)
	{
		if (this->process_data[i] != 0)
		{
			num++;
			sum += this->process_data[i] ^ 2;
		}
	}
	if (num != 0)
	{
		this->average_loudness = int(sqrt(sum / num));
	}
	else
	{
		this->average_loudness = 0;
	}
	LOG("average loudness: ", this->average_loudness, "\n");
	printMemoryInfo("calculate average loudness after");
}

void wavAudio::calculateTime()
{
	LOG("\n\nbegin to calculate time\n");
	this->duration = float(this->end_clock - this->begin_clock) / float(this->sample_rate * this->channels);
	printMemoryInfo("calculate time after");
}

void wavAudio::calculateEnergy()
{
	LOG("\n\nbegin to calculate energy\n");
	double energy = 0;
	float d = 1 / float(this->sample_rate * this->channels);
	for (int i = this->begin_clock; i < this->end_clock; i++)
	{
		if (this->process_data[i] != 0)
		{
			energy += (this->process_data[i] ^ 2) * d;
		}
	}
	this->energy = energy;
	LOG("energy: ", this->energy, "\n");
	printMemoryInfo("calculate energy after");
}

void wavAudio::calculateSTFT()
{
	LOG("\n\nbegin to calculate STFT\n");
	float *x = Malloc(float, this->samples_num);
	printMemoryInfo("malloc x after");
	std::transform(this->process_data, this->process_data + this->samples_num, x,
				   [](int16_t a)
				   {
					   return static_cast<float>(a) / 32767.f;
				   });
	int n_fft = 400;  // 从400调整到256,128
	int n_hop = 160;  // 从160调整到128
	int n_mel = 40;
	int fmin = 80;
	int fmax = 7600;
	LOG("transform data done\n");


	float *x1 = Malloc(float, DATA_NUM);
	printMemoryInfo("malloc x1 after");
	
	int i = 0;
	for (int index = 0; i < DATA_NUM && index < this->samples_num; index++)
	{
		if (x[index] != 0)
		{
			x1[i++] = x[index];
		}
	}
	Free(x);
	LOG("precess STFT data done\n");

	// if (SD.exists("/process_data_x_x1.txt")) {
	// 	SD.remove("/process_data_x_x1.txt");
	// }
	// File outfile = SD.open("/process_data_x_x1.txt", FILE_WRITE);
	// outfile.println("process_data:");
	// for (int i = 0; i < this->samples_num; i++) {
	// 	outfile.print(x[i]);
	// 	outfile.print(" ");
	// }
	// outfile.println();
	// outfile.println("x:");
	// int x_count = 0;
	// for (int i = 0; i < this->samples_num; i++) {
	// 	outfile.print(x[i], 7);
	// 	outfile.print(" ");
	// 	x_count++;
	// }
	// outfile.println();
	// outfile.print("x_count: ");
	// outfile.println(x_count);
	// outfile.println("x1:");
	// int x1_count = 0;
	// for (int i = 0; i < DATA_NUM; i++) {
	// 	outfile.print(x1[i], 7);
	// 	outfile.print(" ");
	// 	x1_count++;
	// }
	// outfile.println();
	// outfile.print("x1_count: ");
	// outfile.println(x1_count);
	// outfile.close();
	
	// x1完全一致，

	float*** X = simplifiedSTFT(x1, i, n_fft, n_hop, "hann", true, "reflect");
	Free(x1);
	printMemoryInfo("simplifiedSTFT after");

	// 分配 real_part 和 imag_part
	int n_frames = 1 + i / n_hop;
	float *real_part = Malloc(float, (n_fft / 2 + 1) * (n_frames));
	float *imag_part = Malloc(float, (n_fft / 2 + 1) * (n_frames));
	if (!real_part || !imag_part) {
		LOG("Failed to allocate memory for real_part or imag_part\n");
		Free(real_part);
		Free(imag_part);
		return;
	}
	LOG("malloc real_part and imag_part done\n");

	int real_index = 0;
	int imag_index = 0;
	LOG("n_frames:", n_frames, "\n");
	LOG("n_fft / 2:", n_fft / 2, "\n");
	for (int row = 0; row < n_frames; ++row) {
		for (int col = 0; col < n_fft / 2 + 1; ++col) {
			// 检查索引是否超出界限
			if (real_index < (n_fft / 2 + 1) * (n_frames - 1) && imag_index < (n_fft / 2 + 1) * (n_frames - 1)) {
				real_part[real_index++] = X[row][col][0];  // 实部
				imag_part[imag_index++] = X[row][col][1];  // 虚部
			} else {
				LOG("Index out of bounds while assigning real and imag parts\n");
				break;
			}
		}
	}

	if (SD.exists("/stft.txt")) {
		SD.remove("/stft.txt");
	}
	// File file = SD.open("/stft.txt", FILE_WRITE);
	// file.println("real_part:");
	// for (int i = 0; i < real_index; ++i) {
	// 	file.print(real_part[i], 7);
	// 	file.print(" ");
	// }
	// file.println();

	// file.println("imag_part:");
	// for (int i = 0; i < imag_index; ++i) {
	// 	file.print(imag_part[i], 7);
	// 	file.print(" ");
	// }
	// file.println();
	// file.print("row: ");
	// file.println(n_frames);
	// file.print("col: ");
	// file.println(n_fft / 2 + 1);
	// file.close();

	LOG("calculate real_part and imag_part done\n");

	// 计算实部的平均值、方差、最大值和最小值
	this->sftf_real_mean = std::accumulate(real_part, real_part + real_index, 0.0f) / real_index;
	float sum_of_squares_real = std::inner_product(real_part, real_part + real_index, real_part, 0.0f);
	this->sftf_real_variance = sum_of_squares_real / real_index;
	this->sftf_real_min = *std::min_element(real_part, real_part + real_index);
	this->sftf_real_max = *std::max_element(real_part, real_part + real_index);

	LOG("sftf_real_mean: ", this->sftf_real_mean, "\n");
	LOG("sftf_real_variance: ", this->sftf_real_variance, "\n");
	LOG("sftf_real_min: ", this->sftf_real_min, "\n");
	LOG("sftf_real_max: ", this->sftf_real_max, "\n");

	// 计算虚部的平均值、方差、最大值和最小值
	this->sftf_imag_mean = std::accumulate(imag_part, imag_part + imag_index, 0.0f) / imag_index;
	float sum_of_squares_imag = std::inner_product(imag_part, imag_part + imag_index, imag_part, 0.0f);
	this->sftf_imag_variance = sum_of_squares_imag / imag_index;
	this->sftf_imag_min = *std::min_element(imag_part, imag_part + imag_index);
	this->sftf_imag_max = *std::max_element(imag_part, imag_part + imag_index);

	LOG("sftf_imag_mean: ", this->sftf_imag_mean, "\n");
	LOG("sftf_imag_variance: ", this->sftf_imag_variance, "\n");
	LOG("sftf_imag_min: ", this->sftf_imag_min, "\n");
	LOG("sftf_imag_max: ", this->sftf_imag_max, "\n");


	// 分配存储 STFT 结果的数组
	this->sftfRealArray = Malloc(float, SFTF_NUM);
	this->sftfImagArray = Malloc(float, SFTF_NUM);
	LOG("malloc sftfRealArray done\n");

	for (int i = 0; i < SFTF_NUM; ++i) {
		this->sftfRealArray[i] = 0;
		this->sftfImagArray[i] = 0;
	}


	int index = 0;
	int count = 0;

	while (count < SFTF_NUM && index < real_index && index < imag_index) {
		if (real_part[index] != 0 && imag_part[index] != 0) {
			this->sftfRealArray[count] = real_part[index];
			this->sftfImagArray[count] = imag_part[index];
			count++;
		}
		index++;
	}
	// for (int j = 0; j < SFTF_NUM; ++j)
	// {
	// 	while (real_part[index] == 0)
	// 	{
	// 		index++;
	// 		if (index >= real_index)
	// 		{
	// 			index--;
	// 			break;
	// 		}
	// 	}
	// 	if (index < real_index)
	// 	{
	// 		this->sftfRealArray[j] = real_part[index];
	// 	}
	// }
	LOG("calculate STFT done\n");;
	
	// 释放所有动态分配的内存
	Free(real_part);
	Free(imag_part);

	printMemoryInfo("calculate STFT after");
}

void wavAudio::calculateMelspectrogram()
{
	LOG("\n\nbegin to calculate Melspectrogram\n");
	float *x = Malloc(float, this->samples_num);
	
	std::transform(this->process_data, this->process_data + this->samples_num, x,
				   [](int16_t a)
				   {
					   return static_cast<float>(a) / 32767.f;
				   });
	LOG("transform data done\n");

	// 计算梅尔频谱图
	int n_fft = 400;
	int n_hop = 160;
	int n_mel = 40;
	int fmin = 80;
	int fmax = 7600;

	float *x1 = Malloc(float, DATA_NUM);
	LOG("malloc x1 done\n");
	
	int count = 0;
	for (int i = 0, index = 0; i < DATA_NUM && index < this->samples_num; index++)
	{
		if (x[index] != 0)
		{
			x1[i++] = x[index];
			count++;
		}
	}
	Free(x);
	LOG("free x done\n");
	LOG("count:", count, "\n");

	float *x1_resized = Malloc(float, count);
	std::copy(x1, x1 + count, x1_resized);
	Free(x1);
	LOG("free x1 done\n");
	LOG("malloc x1_resized done\n");

	float **mels = melSpectrogram(x1_resized, count, sample_rate, n_fft, n_hop, "hann", true, "reflect", 2.f, n_mel, fmin, fmax);

	Free(x1_resized);
	LOG("free x1_resized done\n");

	// 将二维数组打平为一维数组
	int flattenedMels_size = n_mel * (count / n_hop + 1);
	float *flattenedMels = Malloc(float, flattenedMels_size);
	int idx = 0;
	for (int i = 0; i <= count / n_hop; i++)
	{
		for (int j = 0; j < n_mel; j++)
		{
			flattenedMels[idx++] = mels[i][j];
		}
	}
	LOG("flatten mels done\n");

	// 释放 melSpectrogram 返回的二维数组
	for (int i = 0; i < count / n_hop; i++)
	{
		Free(mels[i]);
	}
	Free(mels);
	LOG("free mels done\n");

	if (SD.exists("/mels.txt")) {
		SD.remove("/mels.txt");
	}
	// File file = SD.open("/mels.txt", FILE_WRITE);
	// for (int i = 0; i < flattenedMels_size; i++)
	// {
	// 	file.print(flattenedMels[i], 7);
	// 	file.print(" ");
	// }
	// file.println();
	// file.print(flattenedMels_size);
	// file.println();
	// file.print("row:");
	// file.print(count / n_hop + 1);
	// file.print(" col:");
	// file.print(n_mel);
	// file.close();

	// 计算平均值
	this->mels_mean = std::accumulate(flattenedMels, flattenedMels + flattenedMels_size, 0.0f) / flattenedMels_size;
	LOG("mels_mean: ", this->mels_mean, "\n");

	// 计算方差
	float sum_of_squares = 0.0f;
	for (int i = 0; i < flattenedMels_size; i++) {
		sum_of_squares += (flattenedMels[i] - this->mels_mean) * (flattenedMels[i] - this->mels_mean);
	}
	this->mels_variance = sum_of_squares / flattenedMels_size;
	LOG("mels_variance: ", this->mels_variance, "\n");
	
	// 计算最小值和最大值
    this->mels_min = *std::min_element(flattenedMels, flattenedMels + flattenedMels_size);
    this->mels_max = *std::max_element(flattenedMels, flattenedMels + flattenedMels_size);
	LOG("mels_min: ", this->mels_min, "\n");
	LOG("mels_max: ", this->mels_max, "\n");

	// calculate the top MELS_NUM dimension
	// float *tmp = Malloc(float, MELS_NUM);
	// std::fill(tmp, tmp + MELS_NUM, 0);
	
	this->melsArray = Malloc(float, MELS_NUM);
	for (int i = 0; i < MELS_NUM; i++)
	{
		this->melsArray[i] = 0;
	}

	int index = 0;
	int countMel = 0;
	while (countMel < MELS_NUM && index < flattenedMels_size) {
		if (flattenedMels[index] != 0) {
			this->melsArray[countMel] = flattenedMels[index];
			countMel++;
		}
		index++;
	}
	Free(flattenedMels);

	LOG("calculate Melspectrogram done\n");
}

bool wavAudio::computeCharacter()
{
	LOG("\n\nbegin to calculate character...\n");
	bool res = this->processdata();
	if (!res)
	{
		LOG("process data failed\n");
		return false;
	}

	this->calculatePitch();

	this->calculateAverageLoudness();

	this->calculateTime();

	this->calculateEnergy();

	this->calculateSTFT();

	printMemoryInfo("Memory info after calculate STFT");
	
	this->calculateMelspectrogram();

	printMemoryInfo("Memory info after calculate Melspectrogram");
	
	return true;
}

float* wavAudio::returnCharacter()
{
    // 计算结果数组的大小
    int size = 16 + MELS_NUM + SFTF_NUM * 2;

    // 动态分配内存
    float* res = Malloc(float, size);
    if (!res)
    {
        // 处理内存分配失败
        size = 0;
        return nullptr;
    }

    // 填充数据
    res[0] = this->duration;
    res[1] = this->pitch;
    res[2] = this->average_loudness;
    res[3] = this->energy;
    res[4] = this->mels_mean; // 误差
    res[5] = this->mels_variance; // 误差
    res[6] = this->mels_min; // 误差
    res[7] = this->mels_max; 
    res[8] = this->sftf_real_mean; // 误差
    res[9] = this->sftf_real_variance; // 小误差
    res[10] = this->sftf_real_max;
    res[11] = this->sftf_real_min; // 小误差
    res[12] = this->sftf_imag_mean; // 误差
    res[13] = this->sftf_imag_variance;
    res[14] = this->sftf_imag_max;
    res[15] = this->sftf_imag_min;

    // 填充 melsArray 数据
    for (int i = 0; i < MELS_NUM; ++i)
    {
        res[16 + i] = this->melsArray[i];
    }

    // 填充 sftfRealArray 数据
	int index = 0;
    for (int i = 0; i < SFTF_NUM;)
    {
        res[16 + MELS_NUM + i] = this->sftfRealArray[index];
		res[16 + MELS_NUM + i + 1] = this->sftfImagArray[index];
		index++;
		i = i + 2;
    }

    LOG("return character done\n");
    for (int i = 0; i < size; i++)
    {
        // LOG(i, ":", res[i], " ");
        Serial.print(i);
        Serial.print(":");
        Serial.print(res[i], 7);
        Serial.print(" ");
    }
    LOG("\n");
    // 返回结果
    return res;
}


wavAudio::~wavAudio()
{
	Free(this->data);
	Free(this->process_data);
}
